import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import accuracy_score

# 设置随机种子保证可重复性
torch.manual_seed(42)


# 1. 数据准备
def load_data(dataset_name):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)) if dataset_name == 'FashionMNIST'
        else transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    if dataset_name == 'FashionMNIST':
        train_dataset = torchvision.datasets.FashionMNIST(root='./data',
                                                          train=True,
                                                          transform=transform,
                                                          download=True)
        test_dataset = torchvision.datasets.FashionMNIST(root='./data',
                                                         train=False,
                                                         transform=transform)
        input_size = 28 * 28
        num_classes = 10
    else:  # CIFAR-10
        train_dataset = torchvision.datasets.CIFAR10(root='./data',
                                                     train=True,
                                                     transform=transform,
                                                     download=True)
        test_dataset = torchvision.datasets.CIFAR10(root='./data',
                                                    train=False,
                                                    transform=transform)
        input_size = 32 * 32 * 3
        num_classes = 10

    # 创建数据加载器
    batch_size = 64
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size,
                                              shuffle=False)

    return train_loader, test_loader, input_size, num_classes


# 2. 定义神经网络模型
class ThreeLayerNet(nn.Module):
    def __init__(self, input_size, num_classes):
        super(ThreeLayerNet, self).__init__()
        self.fc1 = nn.Linear(input_size, 256)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(256, 128)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(128, num_classes)

    def forward(self, x):
        if len(x.shape) == 4:  # 对于CIFAR-10 (batch, 3, 32, 32)
            x = x.view(x.size(0), -1)
        else:  # 对于FashionMNIST (batch, 1, 28, 28)
            x = x.view(-1, 28 * 28)
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        x = self.relu2(x)
        x = self.fc3(x)
        return x


# 3. 训练和测试函数
def train_and_test(model, train_loader, test_loader, num_epochs=10, dataset_name='FashionMNIST'):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    train_loss_history = []
    train_acc_history = []
    test_loss_history = []
    test_acc_history = []

    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0

        for images, labels in train_loader:
            images = images.to(device)
            labels = labels.to(device)

            # 前向传播
            outputs = model(images)
            loss = criterion(outputs, labels)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 记录训练数据
            train_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()

        # 计算训练集上的平均loss和准确率
        train_loss = train_loss / len(train_loader)
        train_acc = 100 * train_correct / train_total
        train_loss_history.append(train_loss)
        train_acc_history.append(train_acc)

        # 测试阶段
        model.eval()
        test_loss = 0.0
        test_correct = 0
        test_total = 0

        with torch.no_grad():
            for images, labels in test_loader:
                images = images.to(device)
                labels = labels.to(device)

                outputs = model(images)
                loss = criterion(outputs, labels)

                test_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                test_total += labels.size(0)
                test_correct += (predicted == labels).sum().item()

        # 计算测试集上的平均loss和准确率
        test_loss = test_loss / len(test_loader)
        test_acc = 100 * test_correct / test_total
        test_loss_history.append(test_loss)
        test_acc_history.append(test_acc)

        print(f'Epoch [{epoch + 1}/{num_epochs}], '
              f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%, '
              f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')

    # 可视化训练过程
    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(train_loss_history, label='Train Loss')
    plt.plot(test_loss_history, label='Test Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title(f'{dataset_name} - Training and Testing Loss')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(train_acc_history, label='Train Accuracy')
    plt.plot(test_acc_history, label='Test Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title(f'{dataset_name} - Training and Testing Accuracy')
    plt.legend()

    plt.tight_layout()
    plt.show()

    # 可视化预测结果
    def imshow(img, is_fashionmnist):
        if is_fashionmnist:
            img = img * 0.5 + 0.5  # 反归一化
            npimg = img.numpy()
            plt.imshow(npimg[0], cmap='gray')
        else:
            img = img * 0.5 + 0.5  # 反归一化
            npimg = img.numpy()
            plt.imshow(np.transpose(npimg, (1, 2, 0)))
        plt.axis('off')

    # 获取一批测试数据
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    images = images.to(device)

    # 预测
    outputs = model(images)
    _, predicted = torch.max(outputs, 1)

    # 可视化结果
    images = images.cpu()
    plt.figure(figsize=(12, 8))
    for i in range(12):
        plt.subplot(3, 4, i + 1)
        imshow(images[i], dataset_name == 'FashionMNIST')
        color = 'green' if predicted[i].item() == labels[i].item() else 'red'
        plt.title(f'Pred: {predicted[i].item()}, True: {labels[i].item()}', color=color)
    plt.tight_layout()
    plt.show()

    # 打印最终测试准确率
    print(f'Final Test Accuracy on {dataset_name}: {test_acc_history[-1]:.2f}%')

    return test_acc_history[-1]


# 4. 在FashionMNIST上训练和测试
print("=== Experiment on FashionMNIST ===")
fashion_train, fashion_test, fashion_input_size, fashion_num_classes = load_data('FashionMNIST')
fashion_model = ThreeLayerNet(fashion_input_size, fashion_num_classes)
fashion_acc = train_and_test(fashion_model, fashion_train, fashion_test, dataset_name='FashionMNIST')

# 5. 在CIFAR-10上测试
print("\n=== Experiment on CIFAR-10 ===")
cifar_train, cifar_test, cifar_input_size, cifar_num_classes = load_data('CIFAR10')
cifar_model = ThreeLayerNet(cifar_input_size, cifar_num_classes)
cifar_acc = train_and_test(cifar_model, cifar_train, cifar_test, dataset_name='CIFAR10')

# 6. 结果比较
print("\n=== Final Results ===")
print(f"FashionMNIST Test Accuracy: {fashion_acc:.2f}%")
print(f"CIFAR-10 Test Accuracy: {cifar_acc:.2f}%")
